
import os
os.environ['CUDA_VISIBLE_DEVICES']='0'
import pdb,shutil
import warnings
warnings.filterwarnings('ignore')
# from diffusers import FluxFillPipeline,FluxPriorReduxPipeline,FluxControlPipeline
from diffusers.utils import load_image
from itertools import product
# from image_gen_aux import DepthPreprocessor
from lotus.app_infer_depth import get_depth_by_lotus_g,load_pipe_g

depth_processor = load_pipe_g()

from util_flux import pad_image
# from util_flux import horizontal_concat_images

# from util_sam import get_mask_by_sam
# from util_mask import get_erosed_mask_by_radtio,add_random_holes
import argparse
parser = argparse.ArgumentParser(description="输入提取的类型")
parser.add_argument('-t','--choose_index',required=False,default=0,type=int,help='类型 in [0,1,2]')
args = parser.parse_args()

FLUX_FILL='/home/shengjie/ckp/FLUX.1-Fill-dev'
FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX_DEPTH='/home/shengjie/ckp/FLUX.1-Depth-dev'
FLUX_DEPTH_LORA='/home/shengjie/ckp/FLUX.1-Depth-dev-lora'
FLUX='/data/models/FLUX___1-dev'

DEPTH_PREDCITION='/home/shengjie/ckp/depth-anything-large-hf'


target_shape = (1024,1024)

types = ['collar','sleeve','pockets']
if args.choose_index >= len(types):
    exit(0)
choose_type = types[ int(args.choose_index) ]
examples_dir = f'/mnt/nas/shengjie/datasets/cloth_{choose_type}_balanced/'

depth_dir = f'/mnt/nas/shengjie/datasets/cloth_{choose_type}_balanced_depth/'
# embding_redux_dir = '/data/shengjie/style_zhenzhi_emb/'
if os.path.exists(depth_dir):shutil.rmtree(depth_dir)
os.makedirs(depth_dir)
# os.makedirs(embding_redux_dir,exist_ok=True)

# imagefiles = os.listdir(examples_dir)

# processor = DepthPreprocessor.from_pretrained(DEPTH_PREDCITION)

from tqdm import tqdm
count = 0
for entry in os.scandir(examples_dir):
    filename = entry.name
# for root,subdirs,filenames in os.walk(examples_dir):
    # for filename in tqdm(filenames):
    if not os.path.splitext( filename )[-1] in ['.jpg',]:continue
    count += 1
    if count % 5== 0:
        print(count , end=' ',flush=True) 
    if count % 1000 == 0:
        print()


    img1_path = os.path.join(examples_dir,filename)
    depth_save_path = os.path.join(depth_dir,filename.replace('.jpg','.png'))
    
    # pdb.set_trace()
    # PIL
    img1 = load_image(img1_path) # PIL
    img1,_,_,_,_ = pad_image(img1)
    img1 = img1.resize(target_shape)

    ## depth(img1)
    # img1_depth = processor(img1)[0].convert("RGB")
    # del processor
    # img1_depth.save(depth_save_path)
    output_g = get_depth_by_lotus_g(depth_processor,
                                             img1,None,'cuda') # PIL*2    922*1050
    control_image = output_g.resize(target_shape)

    control_image.save(depth_save_path)

    ## emb(img2)
    # pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
    #                                             FLUX_REDUX, 
    #                                             torch_dtype=torch.bfloat16).to("cuda")
    # main_condition_prompt = pipe_prior_redux(img2) # attr 'prompt_embeds' torch.Size([1, 1241, 4096]) 
    # del pipe_prior_redux